import os
import json


# json_folder = 'E:\BaiduNetdiskDownload\Chip\Chip\lable'  # JSON文件夹路径
# images_dir = 'E:\BaiduNetdiskDownload\Chip\Chip\imges'     # 图像文件夹路径
# output_dir = 'E:\BaiduNetdiskDownload\Chip\Chip\yolo_labels'  # 输出YOLO标注文件夹路径
# 指定JSON文件夹路径和YOLO保存文件夹路径
json_folder = 'E:\BaiduNetdiskDownload\Chip\json_data\lable' # 替换为实际的JSON文件夹路径
yolo_save_folder = 'E:\BaiduNetdiskDownload\Chip\yolo_data\labels'  # 替换为实际保存YOLO txt文件的文件夹

name = ["0", "1"] #标签名
# 如果保存文件夹不存在，创建它
if not os.path.exists(yolo_save_folder):
    os.makedirs(yolo_save_folder)


# 将实例分割数据转换为YOLO格式
def convert_to_yolo_format(json_data, image_width, image_height):
    yolo_data = []
    shapes = json_data['shapes']

    for shape in shapes:
        label = shape['label']
        points = shape['points']

        # 获取分割边界框的x和y的最大最小值
        x_coords = [p[0] for p in points]
        y_coords = [p[1] for p in points]

        x_min = min(x_coords)
        x_max = max(x_coords)
        y_min = min(y_coords)
        y_max = max(y_coords)

        # 计算中心点，宽度和高度（归一化到[0, 1]范围）
        x_center = (x_min + x_max) / 2 / image_width
        y_center = (y_min + y_max) / 2 / image_height
        bbox_width = (x_max - x_min) / image_width
        bbox_height = (y_max - y_min) / image_height

        # 归一化分割点
        normalized_points = [(x / image_width, y / image_height) for x, y in points]

        # 创建YOLO格式的字符串 (格式: class_id x_center y_center width height seg_points)
        if label not in name:
            name.append(label)
        class_id = name.index(label)
        yolo_format = f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height} "

        # 添加分割坐标点
        yolo_format += " ".join([f"{x} {y}" for x, y in normalized_points])
        yolo_data.append(yolo_format)

    return yolo_data


# 遍历文件夹中的所有JSON文件
for filename in os.listdir(json_folder):
    if filename.endswith(".json"):
        json_path = os.path.join(json_folder, filename)
        with open(json_path, 'r', encoding='utf-8') as f:
            json_data = json.load(f)

        # 获取图像大小

        image_width = json_data['imageWidth']
        image_height = json_data['imageHeight']

        # 转换为YOLO格式
        yolo_data = convert_to_yolo_format(json_data, image_width, image_height)

        # 保存为txt文件
        txt_filename = os.path.splitext(filename)[0] + ".txt"
        txt_save_path = os.path.join(yolo_save_folder, txt_filename)

        with open(txt_save_path, 'w', encoding='utf-8') as f:
            for line in yolo_data:
                f.write(line + "\n")

print("转换完成并保存为YOLO格式！")

import torch
print(torch.cuda.is_available())
print(torch.version.cuda)

